#include "Renderer.h"
#include "RayTraceSystem.h"
#include "RenderThreadStorage.h"
#include "Scene.h"
#include "Geometry.h"
#include "VectorMath.h"
#include "Debug.h"
#include <float.h>

using namespace GxLibMath;

namespace RayTracePro
{
	Renderer::Renderer(RayTraceSystem * system)
	{
		this->system = system;
		scene = 0;
		Matrix4::CreateIdentityMatrix(Camera.CameraTransform);
		SetMultiSample(1);
	}

	Renderer::~Renderer()
	{
		for (int i = 0; i<threadStorages.Count(); i++)
		{
			delete threadStorages[i];
		}
	}

	void Renderer::AllocateThreadStorage(int threads)
	{
		for (int i = threadStorages.Count(); i<threads; i++)
			threadStorages.Add(new RenderThreadStorage(0));
		for (int i = 0; i<threadStorages.Count(); i++)
		{
			threadStorages[i]->PhotonSearchBuffer.Photons.SetCapacity(system->PhotonsPerEstimate);
			threadStorages[i]->PhotonSearchBuffer.Distances.SetCapacity(system->PhotonsPerEstimate);
		}
	}

	void Renderer::SetScene(Scene * _scene)
	{
		scene = _scene;
	}
	Scene * Renderer::GetScene()
	{
		return scene;
	}

	void Renderer::SetPerspective(int width, int height, float fovy)
	{
		Camera.SetPerspective(width, height, fovy);
	}

	void Renderer::SetScreenBuffer(RenderBuffer * buffer)
	{
		screenBuffer = buffer;
	}

	RenderBuffer * Renderer::GetScreenBuffer()
	{
		return screenBuffer;
	}

	void Renderer::PrepareRendering()
	{
		if (system->EnableGI)
		{
			DebugTimming::Start("Precomputing irradiance...");
			Array<Importon, true> importons;
			EmitImportons(scene, importons, Camera);
			scene->GetGlobalPhotonMap()->PrecomputeIrradiance(importons, scene, system->PhotonsPerEstimate, scene->GetGlobalPhotonMapRadius());
			DebugTimming::End();
		}
	}

	void Renderer::RenderBlock(int x, int y, int w, int h)
	{
		DifferentialGeometry dgeom;
		SobolGenerator sobolGen(x*321847, y*321847);
		if (multiSample > 1)
		{
			for (int i = 0; i < w; i++)
			{
				for (int j = 0; j < h; j++)
				{
					dgeom.Init();
					int px = x+i;
					int py = y+j;
					Vec4 color;
					color.SetZero();
					RenderPixelMultiSampled(color, dgeom, px, py, sobolGen);
					screenBuffer->SetPixel(px, py, color);
				}
			}
		}
		else
		{
			for (int i = 0; i < w; i++)
			{
				for (int j = 0; j < h; j++)
				{
					dgeom.Init();
					int px = x+i;
					int py = y+j;
					Vec4 color;
					color.SetZero();
					RenderPixel(color, dgeom, px, py);
					screenBuffer->SetPixel(px, py, color);
				}
			}
		}
	}

	void Renderer::RenderPixel(Vec4 & rs, DifferentialGeometry &dgeom, int x, int y)
	{
		Ray r;
		RayDifferential rayDiff;
		Camera.GenerateRay(r, rayDiff, (float)x, (float)y);
		// TraceRay
		float t = -1.0f;
		if (scene->TraceRayCore(dgeom, r, rayDiff, t))
		{
			// Compute transferred ray differential
			TransferRayDifferential(rayDiff, r, dgeom.Normal, t);
			// Shade intersection
			scene->Shade(rs, dgeom, r, rayDiff, t, 0);
		}
		else
		{
			if (scene->Environment)
				scene->Environment->Shade(rs, r.Direction, rayDiff.dDdx, rayDiff.dDdy);
		}
	}

	void Renderer::RenderPixelMultiSampled(Vec4 & rs, DifferentialGeometry &dgeom, int x, int y, SobolGenerator & randGen)
	{
		rs.SetZero();
		Vec4 tmpRs;
		for (int i = 0; i<multiSample; i++)
		{
			tmpRs.SetZero();
			Ray r;
			RayDifferential rayDiff;
			float subSampleX = (float)randGen.NextX()-0.5f;
			float subSampleY = (float)randGen.NextY()-0.5f;
			Camera.GenerateRay(r, rayDiff, (float)x+subSampleX, (float)y+subSampleY);
			// TraceRay
			float t = -1.0f;
			if (scene->TraceRayCore(dgeom, r, rayDiff, t))
			{
				// Compute transferred ray differential
				TransferRayDifferential(rayDiff, r, dgeom.Normal, t);
				// Shade intersection
				scene->Shade(tmpRs, dgeom, r, rayDiff, t, 0);
			}
			else
			{
				if (scene->Environment)
					scene->Environment->Shade(tmpRs, r.Direction, rayDiff.dDdx, rayDiff.dDdy);
			}
			if (tmpRs.x > 1.0f) tmpRs.x = 1.0f;
			if (tmpRs.y > 1.0f) tmpRs.y = 1.0f;
			if (tmpRs.z > 1.0f) tmpRs.z = 1.0f;
			rs += tmpRs;
		}
		rs.x *= invMultiSample;
		rs.y *= invMultiSample;
		rs.z *= invMultiSample;
	}

	RenderThreadScheduler::RenderThreadScheduler(Renderer * renderer, int blockWidth, int blockHeight)
	{
		int screenWidth = renderer->GetScreenBuffer()->Width;
		int screenHeight = renderer->GetScreenBuffer()->Height;
		int xblocks = Ceil(screenWidth / (double)blockWidth);
		int yblocks = Ceil(screenHeight / (double)blockHeight);
		tasks.SetCapacity(xblocks * yblocks);
		for (int i = 0; i < yblocks; i++)
		{
			for (int j = 0; j < xblocks; j++)
			{
				RenderTask task;
				task.x = j * blockWidth;
				task.y = i * blockHeight;
				if (j == xblocks - 1)
				{
					task.w = screenWidth % blockWidth;
					if (task.w == 0)
						task.w = blockWidth;
				}
				else
					task.w = blockWidth;
				if (i == yblocks - 1)
				{
					task.h = screenHeight % blockHeight;
					if (task.h == 0)
						task.h = blockHeight;
				}
				else
					task.h = blockHeight;
				tasks.Add(task);
			}
		}
		exited_threads = 0;
		this->renderer = renderer;
		taskId = 0;
		remainingTasks = tasks.Count();
	}

	RenderThreadScheduler::~RenderThreadScheduler()
	{
	}

	int RenderThreadScheduler::GetNextTask()
	{
		int rs;
		lock.Lock();
		rs = taskId;
		taskId ++;
		lock.Unlock();
		if (rs < tasks.Count())
			return rs;
		else
			return -1;
	}

	void RenderThreadScheduler::Render(int threadCount)
	{
		Cancel();
		DebugTimming::ResetCounters();
		renderer->AllocateThreadStorage(threadCount);
		renderer->PrepareRendering();
		taskId = 0;
		remainingTasks = tasks.Count();
		this->exited_threads = threadCount;
		cancelled = false;
		threads.SetSize(threadCount);
		for (int i = 0; i < threadCount; i++)
		{
			GxWinSystem::Threading::gxThreadParameterizedProc * proc = new GxWinSystem::Threading::gxThreadParameterizedProc(
				this, &RenderThreadScheduler::RenderThread);
			threads[i] = new GxWinSystem::Threading::gxThread(proc, (Object*)(i));
			threads[i]->Start();
		}
	}

	void RenderThreadScheduler::Cancel()
	{
		cancelled = true;
		for (int i = 0; i<threads.Count(); i++)
			WaitForSingleObject(threads[i]->GetHandle(), INFINITE);
	}

	void RenderThreadScheduler::RenderThread(Object * threadIdObj)
	{
		int threadId = (int)(threadIdObj);
		TlsSetValue(renderer->GetSystem()->TlsId, renderer->GetThreadStorage(threadId));
		while (true)
		{
			int taskId = GetNextTask();
			if (taskId == -1)
			{
				break;
			}
			const RenderTask & task = tasks[taskId];
			BlockAssigned(task);
			renderer->RenderBlock(task.x, task.y, task.w, task.h);
			BlockCompleted(task);
			_InterlockedExchangeAdd(&remainingTasks, -1);
			if (remainingTasks == 0)
			{
				RenderCompleted(true);
			}
			else if (cancelled)
			{
				break;
			}
		}
		_InterlockedExchangeAdd(&exited_threads, -1);
		if (exited_threads == 0 && cancelled)
			RenderCompleted(false);
	}


	void Renderer::DebugDrawPhotons()
	{
		PhotonMap * pm = scene->GetGlobalPhotonMap();
		Matrix4 invCamTrans;
		Camera.CameraTransform.Inverse(invCamTrans);
		Vec3 camDir = Vec3(0.0f, 0.0f, 1.0f);
		Vec3 camDir_t;
		Camera.CameraTransform.TransformNormal(camDir_t, camDir);
		for (int i = 1; i<pm->Photons.Count(); i++)
		{
			Ray ray;
			ray.Origin.SetZero();
			ray.Origin.x = Camera.CameraTransform.values[12];
			ray.Origin.y = Camera.CameraTransform.values[13];
			ray.Origin.z = Camera.CameraTransform.values[14];
			Vec3::Subtract(ray.Direction, pm->Photons[i].Position, ray.Origin);
			float maxT = ray.Direction.Length();
			ray.Direction *= 1.0f/maxT;
			ray.ReciprocalDirection.x = 1.0f/ray.Direction.x;
			ray.ReciprocalDirection.y = 1.0f/ray.Direction.y;
			ray.ReciprocalDirection.z = 1.0f/ray.Direction.z;
			ray.tMin = 0.0f;
			ray.tMax = maxT - 0.5f;
			if (!scene->TraceRayPredCore(maxT, ray))
			{
				Vec3 dir;
				invCamTrans.TransformNormal(dir, ray.Direction);
				if (dir.z > -Epsilon)
					continue;
				float scale = Camera.screenZ/dir.z;
				Vec3::Scale(dir, dir, scale);
				int x = (int)dir.x + Camera.centerX;
				int y = -(int)dir.y + Camera.centerY;
				if (x>=0 && x<screenBuffer->Width && y>=0 && y<=screenBuffer->Height)
					screenBuffer->SetPixel(x, y, Vec4(1.0f, 0.0f, 0.0f, 1.0f));
			}
		}
	}
}